import requests
import json
import re

def get_access_token():
    CLIENT_ID = '6XUa10kSl6lDCokWvACoWeKL'
    CLIENT_SECRET = '2Cef6J4fVjG4lKwjoSk1hTC1T1DTrkCT'
    CLIENT_ID1 = 'YrnG3qP8oaFLB9YQvXFcW7eL'
    CLIENT_SECRET1 = 'rp0Ad1QDHGOU5HyA2TARvfphLUlszmza'
    # url = "https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id=&client_secret="
    url = f"https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id={CLIENT_ID1}&client_secret={CLIENT_SECRET1}"
    payload = json.dumps("")
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload)
    return response.json().get("access_token")


def main2(content=None):
    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token=" + get_access_token()

    payload = json.dumps({
        "messages": [
            {
                "role": "user",
                "content": content
            }
        ],
        "stream": True
    })
    headers = {
        'Content-Type': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload, stream=True)

    print(response.text)
    return response.text


def main3():
    # 原始字符串
    data_string = """data: {\"id\":\"as-9ci484qs9h\",\"object\":\"chat.completion\",\"created\":1714966346,\"sentence_id\":0,\"is_end\":false,\"is_truncated\":false,\"result\":\"你好，\",\"need_clear_history\":false,\"finish_reason\":\"normal\",\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\ndata: {\"id\":\"as-9ci484qs9h\",\"object\":\"chat.completion\",\"created\":1714966346,\"sentence_id\":1,\"is_end\":false,\"is_truncated\":false,\"result\":\"很高兴与你交流。\",\"need_clear_history\":false,\"finish_reason\":\"normal\",\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\ndata: {\"id\":\"as-9ci484qs9h\",\"object\":\"chat.completion\",\"created\":1714966346,\"sentence_id\":2,\"is_end\":false,\"is_truncated\":false,\"result\":\"请问有什么我可以帮助你的吗？\",\"need_clear_history\":false,\"finish_reason\":\"normal\",\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\ndata: {\"id\":\"as-9ci484qs9h\",\"object\":\"chat.completion\",\"created\":1714966347,\"sentence_id\":3,\"is_end\":true,\"is_truncated\":false,\"result\":\"\",\"need_clear_history\":false,\"finish_reason\":\"normal\",\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":13,\"total_tokens\":14}}\n\n"""

    # 使用正则表达式匹配每个data后的JSON字符串
    matches = re.findall(r'data: (.*?)\n', data_string, re.DOTALL)

    # 提取result值并存储到列表中
    result_values = []
    for match in matches:
        try:
            # 解析JSON字符串
            data = json.loads(match)
            # 提取result值
            result_values.append(data['result'])
        except json.JSONDecodeError:
            # 如果解析失败，可以记录或忽略错误
            pass

            # 打印结果
    print(result_values)


def main(content):
    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token=" + get_access_token()

    payload = json.dumps({
        "messages": [
            {
                "role": "user",
                "content": content
            }
        ]
    })
    headers = {
        'Content-Type': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload)

    print(response.json().get("result"))

    return response.json().get("result")

# 测试调用api
def test_api():
    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token=" + get_access_token()
    payload = json.dumps({
        "messages": [
            {
                "role": "user",
                "content": "你好"
            }
        ]
    })
    headers = {
        'Content-Type': 'application/json'
    }
    try:
        response = requests.request("POST", url, headers=headers, data=payload)
        response.raise_for_status()  # 检查 HTTP 状态码
        print(response.json())
    except requests.exceptions.RequestException as e:
        print(f"请求失败: {e}")

if __name__ == '__main__':
    main('你好')
